import numpy as np
from tqdm import tqdm

from benchmarks.base_benchmark import BaseBenchmark


class BaseOptimizer:
    optimizer_name = 'base_optimizer'

    def __init__(self, config: dict):
        self.n_run = 0
        self.n_part = 0
        self.show = 0
        self.benchmark = BaseBenchmark({})
        self.step_num = 0

        self.n_group = 0
        self.n_dim = 0
        self.pos_max = 0
        self.pos_min = 0

        self.fe_num = 0

        self.fe_max = 10000
        self.run_flag = True
        self.record_per_fe = 1000
        self.result_cache = []
        self.history_best_x_list = []
        self.init_finish = False

        self.max_v = 0

        self.min_v = 0

        for k, v in config.items():
            setattr(self, k, v)

        self.history_best_x = np.zeros(self.n_dim)
        self.history_best_fit = np.inf
        self.atom_history_best_fits = np.zeros(self.n_part) + np.inf
        self.best_index = 0

        self.xs = np.zeros([self.n_part, self.n_dim])
        self.vs = np.zeros_like(self.xs)

        self.xs = np.random.uniform(self.pos_min, self.pos_max, self.xs.shape)
        self.vs = np.random.uniform(self.min_v, self.max_v, self.xs.shape)

        self.fits = np.zeros(self.n_part, )
        self.p_best = self.xs.copy()
        if self.show:
            self.tqdm_bar = tqdm(total=self.fe_max)
        self.clpso_flag = np.zeros_like(self.fits)
        self.real_clpso_flag = np.zeros_like(self.fits)


    def run(self, *args, **kwargs):
        self.step_num = 0
        while self.step_num < self.n_run:
            self.step_num += 1
            if self.run_flag:
                actions = None
                if self.show:
                    print('群运行中：{}'.format(self.step_num))
                self.run_once(actions=actions)
            else:
                print('fe 超过限制 当前fe:{} 最大fe:{}'.format(self.fe_num, self.fe_max))
                break
        print('finish')

    def run_once(self, actions=None):
        raise NotImplementedError

    def add_check_fe(self):
        if (self.fe_num % self.record_per_fe == 0 or self.fe_num == self.fe_max) and self.fe_num <= self.fe_max:
            if self.show:
                self.tqdm_bar.set_description(f'best:{self.history_best_fit}|process:')
                self.tqdm_bar.update(self.record_per_fe)
            self.data_collect_method()
        if self.fe_num >= self.fe_max:
            # self.data_collect_method()
            self.run_flag = False
        self.fe_num += 1

    def fun(self, x):
        if len(x.shape) == 2:
            ans = []
            for row in x:
                y = self.benchmark.evaluate(np.float64(row))
                self.add_check_fe()
                ans.append(y)
        else:
            ans = self.benchmark.evaluate(np.float64(x))
            self.add_check_fe()
        return np.array(ans)

    def data_collect_method(self):
        fits = self.fits
        mean = np.mean(fits)
        std = np.std(fits)
        best = self.history_best_fit
        data = {
            'fe_num': self.fe_num,
            'mean': mean,
            'best': best,
            'std': std,
        }
        self.result_cache.append(data)

    def best_update(self):
        for i in range(self.n_part):
            new_fit = self.benchmark.evaluate(self.xs[i])
            self.add_check_fe()
            if new_fit < self.history_best_fit:
                self.history_best_fit = new_fit
                self.history_best_x = self.xs[i].copy()
                self.best_index = i
            if new_fit < self.atom_history_best_fits[i]:
                self.atom_history_best_fits[i] = new_fit
                self.p_best[i] = self.xs[i].copy()
                self.history_best_x_list.append((np.sign(self.xs[i].copy()), i))
                self.clpso_flag[i] = 0
                self.real_clpso_flag[i] = 0
            else:
                self.clpso_flag[i] += 1
                self.real_clpso_flag[i] += 1
            # if i == self.best_index:
            #     self.history_best_x_list.append((np.sign(self.xs[i].copy()), 11))

            self.fits[i] = new_fit


    def evaluate(self, x):
        new_fit = self.benchmark.evaluate(x)
        self.add_check_fe()
        return new_fit

    def clip(self):
        self.xs[self.xs > self.pos_max] = self.pos_max
        self.xs[self.xs < self.pos_min] = self.pos_min
        self.vs[self.vs > self.max_v] = self.max_v
        self.vs[self.vs < self.min_v] = self.min_v
